package leetcode.code1172;

import java.util.PriorityQueue;

/**
 * 执行用时：156 ms, 在所有 Java 提交中击败了20.93%的用户
 *
 * 内存消耗：125.4 MB, 在所有 Java 提交中击败了55.81%的用户
 *
 * 通过测试用例：15 / 15
 */
public class DinnerPlates3 {
	int[][] stack;
	int[] stackp;
	int max = 200000;
	PriorityQueue<Integer> notFull;
	PriorityQueue<Integer> notEmpty;
	int capacity;
	int index2;

	public DinnerPlates3(int capacity) {
		this.capacity = capacity;
		stack = new int[max + 1][capacity + 1];
		stackp = new int[max + 1];
		notFull = new PriorityQueue<>((a, b) -> a - b);
		notFull.add(0);
		index2 = 1;
		notEmpty = new PriorityQueue<>((a, b) -> b - a);
	}

	public void push(int val) {
		if (notFull.isEmpty()) {
			notFull.add(index2);
			index2++;
		}
		int index = notFull.peek();
		int size = stackp[index];
		if (size == 0) {
			notEmpty.add(index);
		}
		stack[index][size] = val;
		stackp[index]++;
		if (stackp[index] == capacity) {
			notFull.poll();
		}
	}

	public int pop() {
		int index = -1;
		int size = 0;
		while (!notEmpty.isEmpty()) {
			index = notEmpty.peek();
			if ((size = stackp[index]) == 0) {
				notEmpty.poll();
			} else {
				break;
			}
		}
		if (this.notEmpty.isEmpty()) {
			return -1;
		} else {
			stackp[index]--;
			if (stackp[index] == 0) {
				notEmpty.poll();
			}
			if (size == this.capacity) {
				notFull.add(index);
			}
			return stack[index][size - 1];
		}
	}

	public int popAtStack(int index) {
		int size = stackp[index];
		if (size == 0) {
			return -1;
		} else {
			stackp[index]--;
			if (size == this.capacity)
				notFull.add(index);
			return stack[index][size - 1];
		}
	}

	public static void main(String[] args) {
		test1();
		test2();
		test3();

	}

	public static void test3() {
		DinnerPlates3 D = new DinnerPlates3(2); // 初始化，栈最大容量 capacity = 2
		D.push(1);
		D.push(2);
		D.push(3);
		D.push(4);
		int ans = 0;
		ans = D.pop();
		System.out.println(ans);
		D.push(4);
		ans = D.popAtStack(1);
		System.out.println(ans);
	}

	public static void test2() {
		DinnerPlates3 D = new DinnerPlates3(1); // 初始化，栈最大容量 capacity = 2
		D.push(1);
		D.push(2);
		D.push(3);
		int ans = D.popAtStack(1); // 返回 2。栈的现状为：  4
		System.out.println(ans);
		ans = D.pop(); // 返回 5。栈的现状为： 4
		System.out.println(ans);
		ans = D.pop(); // 返回 4。栈的现状为： 1 3
		System.out.println(ans);

	}

	public static void test1() {
		DinnerPlates3 D = new DinnerPlates3(2); // 初始化，栈最大容量 capacity = 2
		D.push(1);
		D.push(2);
		D.push(3);
		D.push(4);
		D.push(5); // 栈的现状为： 2  4
		int ans = D.popAtStack(0); // 返回 2。栈的现状为：  4
		System.out.println(ans);
		D.push(20); // 栈的现状为： 20 4
		D.push(21); // 栈的现状为： 20 4 21
		ans = D.popAtStack(0); // 返回 20。栈的现状为： 4 21
		System.out.println(ans);
		ans = D.popAtStack(2); // 返回 21。栈的现状为： 4
		System.out.println(ans);
		ans = D.pop(); // 返回 5。栈的现状为： 4
		System.out.println(ans);
		ans = D.pop(); // 返回 4。栈的现状为： 1 3
		System.out.println(ans);
		ans = D.pop(); // 返回 3。栈的现状为： 1
		System.out.println(ans);
		ans = D.pop(); // 返回 1。现在没有栈。
		System.out.println(ans);
		ans = D.pop(); // 返回 -1。仍然没有栈。
		System.out.println(ans);
	}

}
